from openai import OpenAI

PROMPT_NEED_MORE_DATA = (
    "你是一个智能助手。请判断下面的问题，是否需要更多的数据才能给出准确答案。"
    "只回答 'yes' 或 'no'，不要输出其它内容。\n"
    "问题：{question}"
)
PROMPT_ANSWER = (
    "请用简洁、专业的语言回答用户的问题：\n{question}"
)

default_model = "TT-14B-R1-0223"
default_base_url = "http://192.168.210.9:81/v1"

class QuestionAnalyzer:
    def __init__(self, model: str = default_model, api_key: str = "sk-local", base_url: str = default_base_url):
        self.client = OpenAI(
            api_key=api_key,
            base_url=base_url
        )
        self.model = model

    def need_more_data(self, question: str) -> str:
        prompt = PROMPT_NEED_MORE_DATA.format(question=question)
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=3,
            temperature=0
        )
        answer = response.choices[0].message.content.strip().lower()
        if "yes" in answer:
            return "yes"
        elif "no" in answer:
            return "no"
        else:
            return "unknown"

    def generate_answer(self, question: str) -> str:
        prompt = PROMPT_ANSWER.format(question=question)
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=512,
            temperature=0.7
        )
        return response.choices[0].message.content.strip() 